import pydicom
import numpy as np
from PIL import Image
import os


def dcm_to_jpg(dcm_path, jpg_path, output_size=None):
    """
    将DICOM文件转换为JPEG

    参数:
        dcm_path: 输入的DCM文件路径
        jpg_path: 输出的JPG文件路径
        output_size: 可选，输出图像尺寸，例如(512, 512)
    """
    try:
        # 读取DICOM文件
        ds = pydicom.dcmread(dcm_path)

        # 提取像素数据
        pixel_array = ds.pixel_array

        # 获取窗宽和窗位
        if 'WindowWidth' in ds and 'WindowCenter' in ds:
            window_width = ds.WindowWidth
            window_center = ds.WindowCenter
            # 处理多个窗宽窗位的情况（如心超可能没有）
            if isinstance(window_width, pydicom.multival.MultiValue):
                window_width = window_width[0]
                window_center = window_center[0]
        else:
            # 如果没有窗宽窗位，使用像素值的范围
            window_width = pixel_array.max() - pixel_array.min()
            window_center = (pixel_array.max() + pixel_array.min()) / 2

        # 应用窗宽窗位调整
        img_min = window_center - window_width // 2
        img_max = window_center + window_width // 2

        # 将像素值缩放到0-255范围
        pixel_array = np.clip(pixel_array, img_min, img_max)
        pixel_array = ((pixel_array - img_min) / (img_max - img_min) * 255).astype(np.uint8)

        # 转换为PIL图像
        if len(pixel_array.shape) == 3 and pixel_array.shape[2] == 3:
            # 彩色图像 (RGB)
            img = Image.fromarray(pixel_array, 'RGB')
        else:
            # 灰度图像
            img = Image.fromarray(pixel_array, 'L')

        # 调整大小（如果指定）
        if output_size:
            img = img.resize(output_size, Image.Resampling.LANCZOS)

        # 保存为JPEG
        img.save(jpg_path, 'JPEG', quality=95)
        print(f"成功转换: {os.path.basename(dcm_path)} -> {os.path.basename(jpg_path)}")

    except Exception as e:
        print(f"转换失败 {dcm_path}: {str(e)}")


def batch_convert_dcm_to_jpg(input_folder, output_folder, output_size=None):
    """批量转换文件夹中的所有DCM文件"""

    if not os.path.exists(output_folder):
        os.makedirs(output_folder)

    dcm_files = [f for f in os.listdir(input_folder) if f.lower().endswith('.dcm')]

    for dcm_file in dcm_files:
        dcm_path = os.path.join(input_folder, dcm_file)
        jpg_filename = os.path.splitext(dcm_file)[0] + '.jpg'
        jpg_path = os.path.join(output_folder, jpg_filename)

        dcm_to_jpg(dcm_path, jpg_path, output_size)


# 使用方法
if __name__ == "__main__":
    # 安装所需库: pip install pydicom pillow numpy

    # 单个文件转换
    dcm_to_jpg("input/ultrasound.dcm", "output/image.jpg")

    # 批量转换
    batch_convert_dcm_to_jpg("D:\\Data\\xinchao\\dcm\\10.25", "D:\\Data\\xinchao\\jpg\\10.25", output_size=(512, 512))